#include "matrix_operations.h"

void matrix_mul(Matrix a, Matrix b, Matrix result) {
    if (a.cols != b.rows) {
        fprintf(stderr, "Error: Number of columns in A must match number of rows in B.\n");
        return;
    }
    
    // 确保result矩阵已分配内存
    if (result.data == NULL) {
        result.data = (uint8_t*)malloc(a.rows * b.cols * sizeof(uint8_t));
        result.rows = a.rows;
        result.cols = b.cols;
    }

    // 初始化结果矩阵为0
    for (int i = 0; i < a.rows * b.cols; i++) {
        result.data[i] = 0;
    }

    // 执行矩阵乘法 
    for (int i = 0; i < a.rows; i++) {
        for (int j = 0; j < b.cols; j++) {
            for (int k = 0; k < a.cols; k++) {
                result.data[i * b.cols + j] += a.data[i * a.cols + k] * b.data[k * b.cols + j];
            }
        }
    }
}

void mul_matrix(int rowsB, int colsC, int rowsC, Matrix B, Matrix C, Matrix *D) {
    // 1. 维度校验：矩阵乘法要求 "左矩阵列数 = 右矩阵行数"
    //    B 的列数应该等于 C 的行数（即 rowsC）
    if (B.cols != rowsC) {
        fprintf(stderr, "矩阵维度不匹配! B 的列数(%d) 应等于 C 的行数(%d)\n", 
                B.cols, rowsC);
        return;
    }

    // 2. 初始化结果矩阵 D 的维度
    D->rows = rowsB;   // 结果矩阵行数 = B 的行数（1）
    D->cols = colsC;   // 结果矩阵列数 = C 的列数（HIDDEN_LAYER_SIZE）
    // 按需重新分配内存（如果之前大小不对，或首次分配）
    if (D->data != NULL) {
        free(D->data);
    }
    D->data = (uint8_t *)malloc(rowsB * colsC * sizeof(uint8_t));
    if (D->data == NULL) {
        fprintf(stderr, "内存分配失败！\n");
        return;
    }

    // 3. 执行矩阵乘法核心逻辑
    //    公式：D[i][j] = sum_{k=0 to rowsC-1} B[i][k] * C[k][j]
    for (int i = 0; i < rowsB; i++) {         // 遍历结果矩阵的行（B 的行）
        for (int j = 0; j < colsC; j++) {     // 遍历结果矩阵的列（C 的列）
            float sum = 0.0f;
            for (int k = 0; k < rowsC; k++) { // 遍历 "公共维度"（C 的行 / B 的列）
                // B 是 1x784，按行优先存储，所以 B[i][k] = B.data[i * B.cols + k]
                // C 是 784x128，所以 C[k][j] = C.data[k * C.cols + j]
                sum += B.data[i * B.cols + k] * C.data[k * C.cols + j];
            }
            // 结果矩阵 D 按行优先存储，D[i][j] = D.data[i * colsC + j]
            D->data[i * colsC + j] = sum;
        }
    }
}


void print_matrix(Matrix m, int limit) {
    printf("Matrix Content (rows=%d, cols=%d):\n", m.rows, m.cols);
    int count = 0;
    
    for (int i = 0; i < m.rows && count < limit; i++) {
        for (int j = 0; j < m.cols && count < limit; j++) {
            printf("%u ", m.data[i * m.cols + j]);
            count++;
        }
        printf("\n");
    }
    
    if (m.rows * m.cols > limit) {
        printf("... (remaining elements omitted)\n");
    }
}

Matrix allocate_matrix(int rows, int cols) {
    Matrix m;
    m.rows = rows;
    m.cols = cols;
    m.data = (uint8_t*)malloc(rows * cols * sizeof(uint8_t));
    return m;
}
